package com.xiashitech.agent.instrumentation.redis.lettuce.inst;

import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.instrumentation.async.boot.AsyncFutureHolder;
import com.xiashitech.agent.instrumentation.base.DebugUtils;
import com.xiashitech.agent.instrumentation.base.InstPreparedArgs;
import com.xiashitech.agent.instrumentation.base.inst.CommonInstrumentation;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSLettuceCallDepth;
import com.xiashitech.agent.instrumentation.redis.RedisData;
import com.xiashitech.agent.instrumentation.redis.lettuce.XSLettuceAsyncCallback;
import com.xiashitech.agent.instrumentation.redis.lettuce.XSLettuceCompleteAction;
import com.xiashitech.agent.instrumentation.redis.lettuce.XsLettuceSingletons;
import com.xiashitech.agent.utils.UserMatchUtil;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.concurrent.CompletableFuture;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * lettuce异步命令的请求插桩
 * @author zhaoyk
 */
public class XSLettuceAsyncInstrumentation extends CommonInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("io.lettuce.core.api.async.RedisStringAsyncCommands");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("io.lettuce.core.api.async.RedisStringAsyncCommands"));
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        transformMethod(isMethod()
                        .and(named("get"))
                        .and(takesArguments(1))
                        .and(isPublic()),
                "GetDataAdvice", typeTransformer);
    }

    public static class GetDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.This Object commands,
                @Advice.Argument(0) Object key,
                @Advice.Local("preparedArgs") InstPreparedArgs preparedArgs,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            preparedArgs = InstPreparedArgs.prepareOnEnter("XSLettuceAsyncInstrumentation.get", true, XSLettuceCallDepth.class);
            if(!InstPreparedArgs.checkOnEnter(preparedArgs))
                return;

            if(key instanceof byte[] && UserMatchUtil.isLoginToken((byte[])key, preparedArgs.getBaggage())) {
                redisData = RedisData.builder().key((byte[]) key).redisOpetion(Instrumentation.Redis_Opetion_Get).build();
                DebugUtils.print4Debug("XSLettuceAsyncInstrumentation.get", new String((byte[])key));
//                context = preparedArgs.getContext();
//                if(XsLettuceSingletons.asyncInstrumenter().shouldStart(context, redisData)) {
//                    context = XsLettuceSingletons.asyncInstrumenter().start(context, redisData);
//                    scope = context.makeCurrent();
//                }
            } else {
                DebugUtils.print4Debug("XSLettuceAsyncInstrumentation.get", "no logintoken");
            }
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return Object value,
                @Advice.Local("preparedArgs") InstPreparedArgs preparedArgs,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            if(!InstPreparedArgs.checkOnExit(preparedArgs))
                return;
//            if (scope == null)
//                return;
//            scope.close();
            if(redisData != null && value instanceof CompletableFuture) {
//                ((CompletableFuture<?>) value).whenComplete(new XSLettuceCompleteAction(context, redisData));
                AsyncFutureHolder.wait((CompletableFuture<?>) value, new XSLettuceAsyncCallback(XsLettuceSingletons.asyncInstrumenter(), redisData));
            }
        }
    }

}
